import torch
import numpy as np
import torch.utils
from torch.utils.data import Dataset, DataLoader
import os
from typing import Dict
import torch.utils.data
import imageio as iio


class OGMImageDataset(Dataset):
    def __init__(self, base_path):
        '''
        NOTE: base_path should contain split name
        '''
        self.base_path = base_path
        self.name, self.split = self.base_path.split('/')[-2:]
        map_files = os.listdir(base_path)
        self.map_files = [os.path.join(base_path, i) for i in map_files]
        self.length = len(self.map_files)
        print(self)
        
    def __len__(self) -> int:
        return self.length
    
    def __getitem__(self, idx) -> Dict[str, torch.tensor]:
        image = iio.imread(self.map_files[idx]).astype(np.float32) / 255.
        image_tensor = torch.from_numpy(image).unsqueeze(0)
        
        data = {
            'maps': image_tensor,
        }
        return data
    
    def __str__(self) -> str:
        return f'Dataset {self.name}, split {self.split}, contains {self.length} maps'
